import os,random,pickle
from tools import *
'''
分类原则是从对应的文本中读取随机生成的编号内容
每个类别训练集选取2000条，测试集选取1000条
要求训练集和测试集互不相交
'''

def make_train_test(path):
    f_train =open('train.txt',mode='w',encoding='utf-8')
    f_test = open('test.txt',mode='w',encoding='utf-8')
    class_dir = os.listdir(path)  #####[13个txt]
    for file in class_dir:
        file_path = path+'/'+file
        ##获取txt的总行数
        n = get_txt_linecount(file_path)
        ##得到全局列表
        all_list = [x for x in range(1,n+1)]
        ##得到随机train表
        train_list=random.sample(range(1,n+1),2000)
        ##从总体中把train去掉
        remain_list=list(set(all_list)^set(train_list))
        ##然后从剩余的行数中选择
        test_list=random.sample(remain_list,1000)
        for t in train_list:
            train_line_content = get_line_content(file_path,t)
            f_train.write(train_line_content)
        for t2 in test_list:
            test_line_content = get_line_content(file_path,t2)
            f_test.write(test_line_content)
    f_train.close()
    f_test.close()

def make_train(path):
    f_train =open('train.txt',mode='w',encoding='utf-8')
    class_dir = os.listdir(path)  #####[13个txt]
    all__=[]
    train__=[]
    for file in class_dir:
        file_path = path+'/'+file
        ##获取txt的总行数
        n = get_txt_linecount(file_path)
        ##得到全局列表
        all_list = [x for x in range(1,n+1)]
        all__.append(all_list) 
        ##得到随机train表
        train_list=random.sample(range(1,n+1),2000)
        train__.append(train_list)
        for tt in train_list:
            train_line_content = get_line_content(file_path,tt)
            f_train.write(train_line_content)
    all_con = [all__,train__,class_dir]
    with open("test.pkl", "wb") as f:
        pickle.dump(all_con, f)
    f.close()
    f_train.close()
        

def make_test():
    f_test = open('test.txt',mode='w',encoding='utf-8')
    fr = open("test.pkl",'rb')# open的参数是pkl文件的路径
    con = pickle.load(fr)  # 读取pkl文件的内容
    fr.close()
    all__=con[0]
    train__=con[1]
    for i in range(13):
        ##从总体中把train去掉
        remain_list=list(set(all__[i])^set(train__[i]))
        ##然后从剩余的行数中选择
        test_list=random.sample(remain_list,1000)
        for t2 in test_list:
            test_line_content = get_line_content('split_data/'+con[2][i],t2)
            f_test.write(test_line_content)
    f_test.close()